package com.example.security;

import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Bucket4j;
import io.github.bucket4j.Refill;
import org.springframework.context.annotation.Configuration;

import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Configuration
public class RateLimitConfig {
    private final Map<String, Bucket> buckets = new ConcurrentHashMap<>();

    public Bucket resolveBucket(String key) {
        return buckets.computeIfAbsent(key, this::newBucket);
    }

    /**
     * 创建一个新的Bucket对象
     *
     * @param key 桶的唯一标识符
     * @return 一个新的Bucket对象
     */
    private Bucket newBucket(String key) {
        // 允许每分钟30次请求
        // 设置带宽限制
        // 每分钟允许100次请求，每1分钟补充10个令牌
        Bandwidth limit = Bandwidth.classic(100, Refill.intervally(10, Duration.ofMinutes(1)));
        // 使用Bucket4j构建器添加带宽限制并构建Bucket对象
        return Bucket4j.builder().addLimit(limit).build();
    }

}